import re

# import urllib.request
# url = ("https://raw.githubusercontent.com/rasbt/"
#        "LLMs-from-scratch/main/ch02/01_main-chapter-code/"
#        "the-verdict.txt")
# file_path = "the-verdict.txt"
# urllib.request.urlretrieve(url, file_path)


# 1.加载文本，训练集
with open("../data/the-verdict.txt", "r", encoding="utf-8") as f:
    raw_text = f.read()
print("Total number os character:", len(raw_text))
print(raw_text[:99])

# 正则使用
# text = "Hello, world. Is this-- a test?"
# result = re.split(r'([,.:;?_!"()\']|--|\s)', text)
# result = [item.strip() for item in result if item.strip()]
# print(result)

# 2.文本分词
preprocessed = re.split(r'([,.:;?_!"()\']|--|\s)', raw_text)
preprocessed = [item.strip() for item in preprocessed if item.strip()]
# print(len(preprocessed))
# print(preprocessed[:50])


# 3.将词元转换为词元ID
# 3.1构建词汇表
all_words = sorted(set(preprocessed))
all_words.extend(["<|endoftext|>", "<|unk|>"])
# vocab_size = len(all_words)
# print(vocab_size)

vocab = {token: idx for idx, token in enumerate(all_words)}
print(type(vocab))  # dict
# for i, item in enumerate(vocab.items()):
for i, item in enumerate(list(vocab.items())[-5:]):
    print(item)
    # if i > 50:
    #     break


class SimpleTokenizerV1:
    def __init__(self, vocab_1):
        self.str_to_int = vocab_1
        self.int_to_str = {i: s for s, i in vocab_1.items()}

    def encode(self, text):
        preprocessed_1 = re.split(r'([,.:;?_!"()\']|--|\s)', text)
        preprocessed_1 = [
            item.strip() for item in preprocessed_1 if item.strip()
        ]
        ids = [self.str_to_int[s] for s in preprocessed_1]
        return ids

    def decode(self, ids):
        text = " ".join([self.int_to_str[i] for i in ids])
        text = re.sub(r'\s+([,.?!"()\'])', r'\1', text)
        return text


class SimpleTokenizerV2:
    def __init__(self, vocab_2):
        self.str_to_int = vocab_2
        self.int_to_str = {i: s for s, i in vocab_2.items()}

    def encode(self, text):
        preprocessed_2 = re.split(r'([,.:;?_!"()\']|--|\s)', text)
        preprocessed_2 = [
            item.strip() for item in preprocessed_2 if item.strip()
        ]
        preprocessed_2 = [item if item in self.str_to_int else "<|unk|>" for item in preprocessed_2]
        ids = [self.str_to_int[s] for s in preprocessed_2]
        return ids

    def decode(self, ids):
        text = " ".join([self.int_to_str[i] for i in ids])
        text = re.sub(r'\s+([,.?!"()\'])', r'\1', text)
        return text



tokenizer = SimpleTokenizerV1(vocab)
text = """"It's the last he painted, you know,"
       Mrs. Gisburn said with pardonable pride."""
ids = tokenizer.encode(text)
print(ids)
# [1, 56, 2, 850, 988, 602, 533, 746, 5, 1126, 596, 5, 1, 67, 7, 38, 851, 1108, 754, 793, 7]
print(tokenizer.decode(ids))
# " It' s the last he painted, you know," Mrs. Gisburn said with pardonable pride.

text = "Hello, do you like tea?"
# print(tokenizer.encode(text))
# KeyError: 'Hello'


tokenizer = SimpleTokenizerV2(vocab)
text1 = "Hello, do you like tea?"
text2 = "In the sunlit terraces of the palace."
text = " <|endoftext|> ".join((text1, text2))
print(text)
# Hello, do you like tea? <|endoftext|> In the sunlit terraces of the palace.
print(tokenizer.encode(text))
# [1131, 5, 355, 1126, 628, 975, 10, 1130, 55, 988, 956, 984, 722, 988, 1131, 7]
print(tokenizer.decode(tokenizer.encode(text)))
# <|unk|>, do you like tea? <|endoftext|> In the sunlit terraces of the <|unk|>.
